import json
import time
from typing import Dict, Any, Optional, List
import socket

class IntelligentRouter:
    """使用AI模型进行路由决策的智能路由器"""
    def __init__(self):
        # 路由决策模型的配置（可以是一个小型高效的模型）
        self.router_model_config = {
            "host": "router-model.example.com",
            "port": 12340,
            "protocol": "mcp-v1-light"
        }
        
        # 可用模型终端的配置
        self.mcp_configs = {
            "deepseek-chat": {
                "host": "deepseek-chat.example.com",
                "port": 12345,
                "protocol": "mcp-v1",
                "description": "通用聊天对话模型"
            },
            "deepseek-code": {
                "host": "deepseek-code.example.com",
                "port": 12346,
                "protocol": "mcp-v1",
                "description": "代码生成与分析模型"
            },
            "deepseek-math": {
                "host": "deepseek-math.example.com",
                "port": 12347,
                "protocol": "mcp-v1",
                "description": "数学问题求解模型"
            },
            "fallback": {
                "host": "default-model.example.com",
                "port": 12348,
                "protocol": "mcp-v2",
                "description": "通用后备模型"
            }
        }
    
    def decide_route(self, user_input: str) -> Dict[str, Any]:
        """使用AI模型决策最佳路由"""
        # 构造路由决策请求
        router_request = {
            "version": "1.0",
            "model": "router-001",
            "messages": [{
                "role": "system",
                "content": ("分析以下用户输入，从以下模型中选择最适合处理的模型:\n"
                           f"可用模型: {json.dumps(list(self.mcp_configs.keys()))}\n"
                           "只需返回最匹配的模型名称，不需要解释。")
            }, {
                "role": "user",
                "content": user_input
            }],
            "parameters": {
                "temperature": 0.2,  # 低随机性，确保稳定路由
                "max_tokens": 10     # 只需要返回模型名称
            }
        }
        
        # 发送路由决策请求
        with MCPClient(**self.router_model_config) as router_client:
            response = router_client.send_request(router_request)
            
        # 解析路由决策结果
        if response and "choices" in response:
            chosen_model = response["choices"][0]["message"]["content"].strip()
            # 清理可能的json字符或多余引号
            chosen_model = chosen_model.replace('"', '').replace("'", "").strip()
            # 检查选择的模型是否在配置中
            if chosen_model in self.mcp_configs:
                return self.mcp_configs[chosen_model]
        
        # 默认返回fallback配置
        print(f"[路由决策] 使用后备模型，原选择: {response.get('choices',[{}])[0].get('message',{}).get('content','')}")
        return self.mcp_configs["fallback"]

class SmartAIClient:
    """集成智能路由的AI客户端"""
    def __init__(self):
        self.router = IntelligentRouter()
        # 对话历史记录
        self.conversation_history = []
    
    def process_input(self, user_input: str) -> str:
        """处理用户输入的全流程"""
        # 添加到对话历史
        self.conversation_history.append({"role": "user", "content": user_input})
        
        # 1. 智能路由选择 - 由AI模型决定
        mcp_config = self.router.decide_route(user_input)
        print(f"[路由选择] 使用模型: {mcp_config['description']}")
        
        # 2. 构造请求（包含对话上下文）
        request = self._build_request(user_input, mcp_config["protocol"])
        
        # 3. 与目标模型交互
        start_time = time.time()
        with MCPClient(mcp_config["host"], mcp_config["port"]) as client:
            response = client.send_request(request)
            
        elapsed = time.time() - start_time
        print(f"[性能统计] 模型处理时间: {elapsed:.2f}s")
        
        # 4. 解析结果
        result = self._parse_response(response)
        
        # 记录AI响应
        self.conversation_history.append({"role": "assistant", "content": result})
        
        return result
    
    def _build_request(self, user_input: str, protocol: str) -> Dict[str, Any]:
        """根据协议构建请求，包含对话上下文"""
        if protocol == "mcp-v1":
            return {
                "version": "1.0",
                "model": "default",
                "messages": self._get_recent_history(5),  # 获取最近的5轮对话
                "parameters": {
                    "temperature": 0.7,
                    "max_tokens": 2000,
                    "top_p": 0.9
                }
            }
        else:  # mcp-v2
            return {
                "version": "2.0",
                "prompt": self._format_history_as_prompt(),
                "config": {
                    "sampling": {
                        "temperature": 0.7,
                        "top_k": 50
                    }
                }
            }
    
    def _get_recent_history(self, turns: int = 3) -> List[Dict[str, str]]:
        """获取最近的对话历史"""
        return self.conversation_history[-2*turns:]  # 获取最近的turns轮对话
    
    def _format_history_as_prompt(self) -> str:
        """将对话历史格式化为mcp-v2需要的提示格式"""
        return "\n".join(
            f"{msg['role']}: {msg['content']}" 
            for msg in self._get_recent_history(3)
        )
    
    def _parse_response(self, response: Optional[Dict[str, Any]]) -> str:
        """解析不同协议的响应"""
        if not response:
            return "抱歉，我遇到了技术问题，暂时无法回答。"
        
        # mcp-v1 格式
        if "choices" in response and len(response["choices"]) > 0:
            return response["choices"][0]["message"]["content"]
        
        # mcp-v2 格式
        if "output" in response:
            return response["output"]
        
        # 其他格式
        return str(response)

class MCPClient:
    """MCP客户端实现（与前例相同）"""
    def __init__(self, host: str, port: int):
        self.host = host
        self.port = port
        self.socket = None
        self.buffer_size = 4096
        self.timeout = 30
    
    def connect(self) -> bool:
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(self.timeout)
            self.socket.connect((self.host, self.port))
            return True
        except Exception as e:
            print(f"[连接错误] {self.host}:{self.port} - {str(e)}")
            return False
    
    def send_request(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        if not self.socket and not self.connect():
            return None
        
        try:
            request_str = json.dumps(request, ensure_ascii=False) + '\n'
            self.socket.sendall(request_str.encode('utf-8'))
            return self._receive_response()
        except socket.timeout:
            print("[错误] 请求超时")
            return None
        except Exception as e:
            print(f"[通信错误] {str(e)}")
            return None
    
    def _receive_response(self) -> Optional[Dict[str, Any]]:
        response_data = b''
        while True:
            try:
                chunk = self.socket.recv(self.buffer_size)
                if not chunk:
                    break
                response_data += chunk
                if b'\n' in chunk:
                    break
            except socket.timeout:
                print("[警告] 接收数据超时")
                break
        
        if not response_data:
            return None
        
        try:
            return json.loads(response_data.decode('utf-8').strip())
        except json.JSONDecodeError as e:
            print(f"[解析错误] {str(e)}")
            return None
    
    def close(self):
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
            finally:
                self.socket = None
    
    def __enter__(self):
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

def main():
    """命令行交互界面"""
    client = SmartAIClient()
    print("智能AI对话系统 (输入'exit'退出)\n")
    
    while True:
        try:
            user_input = input("你: ")
            if user_input.lower() in ('exit', 'quit'):
                break
                
            if not user_input.strip():
                continue
                
            print("\n思考中...\n")
            
            # 处理用户输入
            response = client.process_input(user_input)
            
            print(f"\nAI: {response}\n")
            
        except KeyboardInterrupt:
            print("\n再见！")
            break
        except Exception as e:
            print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main()
